__doc__ =   """ Misc mathematical and geometric functions.
            """

from numpy import mean, sum, arccos, power, dot, linalg, array, finfo, cross, arctan2, c_, vstack, ones, sqrt, sin, cos, pi
#from math import
from random import random
from itertools import izip

def sphere_points(n):
    """ Uniform points on the sphere are generated using the Golden Section Spiral algorithm.
        Boscoh ASA.py (s/cite needed)
    """
    points = []
    inc = pi * (3 - sqrt(5))
    offset = 2 / float(n)
    for k in xrange(int(n)):
        y = k * offset - 1 + (offset / 2)
        r = sqrt(1 - y*y)
        phi = k * inc
        points.append([cos(phi)*r, y, sin(phi)*r])
    return array(points)

def coords_to_mcoords(coords):
    """ Converts an array of coordinates into an array of mcoordinates, for a definition of mcoords see:
        Entity.get_mcoords
    """
    # this version operates on array columns
    coords2 = power(coords, 2)
    mcoords  = array([arctan2(sqrt(coords2[:,1] + coords2[:,2]), coords[:,0]),\
                      arctan2(sqrt(coords2[:,0] + coords2[:,2]), coords[:,1]),\
                      arctan2(sqrt(coords2[:,0] + coords2[:,1]), coords[:,2])
                      ]).transpose()
    return mcoords

def get_rotation_matrix(coords_1, coords_2):
    assert len(coords_1) == len(coords_2) > 0
    coords_1 = coords_1 - mean(coords_1, axis =0)   # centers matrices to
    coords_2 = coords_2 - mean(coords_2, axis =0)   # avoid affine transformation
    # Initial residual, see Kabsch.
    E0 = sum(power(coords_1, 2)) + sum(power(coords_2, 2))
    V, S, Wt = linalg.svd(dot(coords_2.transpose(), coords_1))
    reflect = float(str(float(linalg.det(V) * linalg.det(Wt))))
    if reflect == -1.0:
            S[-1] = -S[-1]
            V[:,-1] = -V[:,-1]
    RMSD = E0 - (2.0 * sum(S))
    RMSD = sqrt(abs(RMSD / len(coords_1)))
    U = dot(V, Wt)               # this is the rotation matrix
    return (U, RMSD)

def random_rotation_matrix(rand1 =None, rand2 =None, rand3 =None):
    """
    Random uniform rotation matrix, based on:
    http://tog.acm.org/GraphicsGems/gemsiii/rand_rotation.c
    """
    rand1 = (rand1 or random())
    rand2 = (rand2 or random())
    rand3 = (rand3 or random())

    theta = rand1 * pi * 2
    phi = rand2 * pi * 2.0
    z = rand3 * 2.0
    r  = sqrt( z );
    Vx = sin( phi ) * r;
    Vy = cos( phi ) * r;
    Vz = sqrt( 2.0 - z );
    st = sin( theta )
    ct = cos( theta )
    Sx = Vx * ct - Vy * st
    Sy = Vx * st + Vy * ct
    M = array ([[Vx * Sx - ct,  Vx * Sy - st,   Vx * Vz],
                [Vy * Sx + st,  Vy * Sy - ct,   Vy * Vz],
                [Vz * Sx,       Vz * Sy,        1.0 - z]])

    return M

def euler_from_rotation_matrix(matrix):
    i, j, k = 0, 1, 2
    M = matrix
    zero = finfo(float).eps * 4.0
    cy = sqrt(M[i,i]*M[i,i] + M[j,i]*M[j,i])
    if cy > zero:
        ax = arctan2( M[k,j],  M[k,k])
        ay = arctan2(-M[k,i],  cy)
        az = arctan2( M[j,i],  M[i,i])
    else:
        ax = arctan2(-M[j,k],  M[j,j])
        ay = arctan2(-M[k,i],  cy)
        az = 0.0
    return ax, ay, az

def kappa_from_rotation_matrix(matrix):
    i, j, k = 0, 1, 2
    M = matrix
    trace = M[i,i] + M[j,j] + M[k,k]
    kappa = arccos((trace - 1.0) / 2.0)
    return kappa

def transform_coords(coords, mx, vec =None, fmx =None, omx =None):
    """ Manipulate coordinates of the asymmetric unit by a rotation matric and translation vector.
    """
    if (fmx != None and omx != None):               #
        coords = dot(coords, fmx.transpose())
    if mx.shape == (4,4):                           #
        coords4 = c_[coords, array([ones(len(coords))]).transpose()]
    rot_mx = mx.transpose()
    new_coords = dot(coords4, rot_mx)[:,:3]                     # rotate / translate
    if (fmx != None and omx != None):
        new_coords = dot(new_coords, omx.transpose())
    # the rotation matrix is stored in the PDB or retrieved from cctbx in a form that requires
    # it's transposition before multiplication with coordinates.
    if vec:
        new_coords + vec
    return new_coords

def coords_to_symmetry(coords, fmx, omx, mxs, mode):
    """ Applies symmetry matrices on coordinates to deerive a P1 unit cell, requires
        orthogonal coordinates, a fractionalization matrix, a orthogonalization matrix and
        rotation matrices.
    """
    all_coords = [coords]                                           # the first matrix is identity
    if mode == 'table':
        coords = dot(coords, fmx.transpose())
    coords4 = c_[coords, array([ones(len(coords))]).transpose()]    # add column of 1.
    for i in xrange(1,len(mxs)):                                    # we skip identity
        rot_mx = mxs[i].transpose()
        new_coords = dot(coords4, rot_mx)[:,:3]                     # rotate / translate
        if mode == 'table':                                                 # no 1. column
            new_coords = dot(new_coords, omx.transpose())
        all_coords.append(new_coords)
    return vstack(all_coords).reshape((len(all_coords),coords.shape[0],3))
    #return array(all_coords) is slower

def coords_to_crystal(coords, fmx, omx, n =1):
    """ Applies primitive lattice translations to produce a crystal from
        the contents of the asymmetric unit with applied symmetries.
    """
    rng = range(-n, n+1)                            # a range like -2, -1, 0, 1, 2
    fcoords = dot(coords, fmx.transpose())          # fractionalize coordinates
    vectors = [(x, y, z) for x in rng for y in rng for z in rng]    # all combinations of the above
    # looking for the center Thickened cube numbers: a(n)=n*(n^2+(n-1)^2)+(n-1)*2*n*(n-1) ;)
    all_coords = []
    for primitive_vector in vectors:
        all_coords.append(fcoords + primitive_vector)
    #all_coords = array(all_coords), the next line is equivalent but faster
    all_coords = vstack(all_coords).reshape((len(all_coords), coords.shape[0], coords.shape[1], 3))
    all_coords = dot(all_coords, omx.transpose())    # back to orthogonal
    return all_coords

def convert_symop(rot_mx, matrix1):
    """ Convert fractional <-> orthogonal symop 4x4 matrix
    """
    matrix2 = linalg.inv(matrix1)                                   # get inverse matrix
    matrix14 = c_[matrix1, array([zeros(3)]).transpose()]           # extend with 0. column yields a 4x3
    matrix24 = c_[matrix2, array([zeros(3)]).transpose()]           # extend with 0. column
    matrix24 = r_[matrix24, array([[0., 0., 0., 1.]])]              # append row yields a 4x4
    # see mmLib code SpaceGroups.py
    tmp_mx = dot(rot_mx.transpose(), matrix14.transpose())
    tmp_mx = dot(matrix24.transpose(), tmp_mx).transpose()
    new_rot_mx = r_[tmp_mx, array([[0., 0., 0., 1.]])]
    return new_rot_mx

#END